[백준] 18352번 - 특정 거리의 도시 찾기 (파이썬)

문제

문제 링크

풀이

처음에는 저번에 다익스트라 알고리즘을 공부하며 썼던 포스팅에서 썼던 예시 코드를 그대로 적용해서 풀어봤는데 시간초과로 실패했다 ㅠ

생각해보니까 예시 코드의 상황과 달리

  • 각 도시는 단방향으로만 연결된다
  • 각 도시 사이의 모든 거리는 1이다

라는 조건이 있는데 이를 고려하지 않고 쓸모없는 계산들을 많이 해서 그런 것 같다.

결국 갓구글의 도움을 받아 완성한 코드..

1) 다익스트라 알고리즘 사용

💡 다익스트라 알고리즘 복습!

  1. 출발 노드를 설정한다
  2. 비용 리스트를 (무한으로) 초기화한다
  3. 방문하지 않은 노드 중 최단 거리가 가장 짧은 노드를 선택한다
  4. 해당 노드를 거쳐 다음 노드로 가는 비용을 계산한다. 다음 노드로 가는 비용을 기존에 구해놓은 값이 있다면, 새로 구한 값과 기존값중 최소값으로 비용 리스트를 갱신한다
  5. 3~4번 과정을 반복한다

💡 다익스트라 알고리즘에서 heap을 사용하는 이유

힙을 이용하지 않고 일반 리스트를 이용한다면 매번 최단 거리가 짧은 노드를 선형 탐색해야 하기 때문에 시간이 오래 걸린다.

힙 자료구조를 이용하면 최솟값이 Root에 오는 특징을 이용해 출발 노드로부터 가장 거리가 짧은 노드를 더 빠르게 찾을 수 있다!

import sys
import heapq

N, M, K, X = map(int, sys.stdin.readline().strip().split(" "))

graph = [[] for _ in range(N+1)]
values = [sys.maxsize]*(N+1)  # 거리배열 int의 최댓값으로 초기화

for i in range(M):
    A, B = map(int, sys.stdin.readline().strip().split(" "))
    graph[A].append(B)  # 모든 거리가 1이기 때문에 거리를 따로 저장할 필요 없이 이렇게만 저장해준다


heap = []  # 거리가 가장 작은 것이 root에 올 수 있게 최소힙 만듬 (출발도시에서 가까운 도시부터 탐색하기 위해)
heapq.heappush(heap, (0, X))  # (해당 도시까지 가는 거리, 도시 번호)
values[X] = 0

while heap:
    v_now, now = heapq.heappop(heap)
    if v_now > values[now]:  # 이미 now까지 가는 최단거리가 구해져 있다면 굳이 계산할 필요 없음. 아래 로직 무시
        continue
    for next in graph[now]:  
        v = v_now + 1
        if v < values[next]: # 새로 구한 값이 더 최단거리이면
            values[next] = v  # 새로 구한 값으로 거리배열 갱신
            heapq.heappush(heap, (v, next))

cnt = 0
for i in range(len(values)):
    if values[i] == K:
        cnt += 1
        print(i)
if cnt == 0:
    print(-1)

2) BFS(너비우선탐색) 알고리즘 사용

⚡ 모든 거리가 동일하기 때문에 BFS로도 풀이가 가능하다

(출발노드로부터 특정 레벨의 노드까지의 거리보다 그 다음 레벨의 노드까지의 거리가 무조건 큼을 보장할 수 있기 때문)

import sys
from collections import deque

N, M, K, X = map(int, sys.stdin.readline().strip().split(" "))

graph = [ [] for _ in range(N+1)]
values = [-1]*(N+1)  # 거리배열 -1(아직 방문하지 않았음을 의미)로 초기화
values[X] = 0

for i in range(M):
   A, B = map(int, sys.stdin.readline().strip().split(" "))
   graph[A].append(B)

queue = deque([X])
while queue :
    now = queue.popleft()  # queue에서 앞에서부터 꺼냄 = 같은 레벨의 노드 모두 탐색 후 다음 레벨로 이동 (BFS)
    for next in graph[now]:
        if values[next] == -1:  # next에 아직 방문하지 않은 경우에만 거리배열 갱신 (이미 방문 했다면 이전 레벨에서 구한 값이 무조건 더 작을 것이므로)
            values[next] = values[now]+1  
            queue.append(next)
    # for문을 다 돌고 나면 다음 레벨로 이동

cnt = 0
for i in range(len(values)):
    if values[i] == K:
        cnt += 1
        print(i)
if cnt == 0:
    print(-1)

Written by@[hyem]
Hyem's Dev Note

GitHub